# -*- coding: utf-8 -*-

import torch
import torch.nn as nn
import numpy as np


class QMixNet(nn.Module):

    def __init__(self):
        super(QMixNet, self).__init__()

        """ Agent (o,u) -> MLP -> GRU -> MLP -> Q_a """

    @staticmethod
    def _process_observation(obs, act):
        obs = obs.float()
        act = act.float()
        obs = obs.reshape(-1, 13, 13, 41)
        obs = obs.permute(0, 3, 1, 2)
        return obs, act

    def _compute_agent_q(self, obs, hidden):

        # MLP
        # e = torch.cat((e, act), 1)
        e = self.mlp(e)

        # GRU
        e = e.reshape(-1, 1, self.obs_embedding_dim)
        origin_e, _ = self.gru(e)
        e = e.reshape(-1, self.obs_embedding_dim)

        # MLP

        return e

    def forward(self, obs, act):

        # Part 1: Reshape, convert ...
        obs, act = self._process_observation(obs, act)

        # Part 2: Observation encoding

        origin_e = self._compute_agent_q(obs, act)

        origin_e = origin_e.reshape(-1, 1, self.obs_embedding_dim)

        origin_e, _ = self.gru(origin_e)

        return origin_e


if __name__ == '__main__':
    test_ts = np.random.randn(8, 13, 13, 41)
    test_ts = torch.from_numpy(test_ts).float()

    test_action = np.random.randn(8, 1)
    test_action = torch.from_numpy(test_action).float()

    model = QMixNet()
    out = model(test_ts, test_action)
    print(out.shape)

